import cv2
import numpy as np
import matplotlib.pyplot as plt

# 生成运动模糊核
def make_PSE(kernel_size=15, angle=60):
    PSE = np.diag(np.ones(kernel_size))  # 初始模糊核方向为 -45°
    angle += 45  # 抵消 -45°的影响
    M = cv2.getRotationMatrix2D((kernel_size / 2, kernel_size / 2), angle, 1)
    PSE = cv2.warpAffine(PSE, M, (kernel_size, kernel_size))
    PSE = PSE / PSE.sum()  # 权重归一化
    return PSE

# 在空域对图像进行运动模糊
def motion_blur(image, PSE):
    blurred = cv2.filter2D(image, -1, PSE)  # 2D卷积实现模糊
    cv2.normalize(blurred, blurred, 0, 255, cv2.NORM_MINMAX)
    blurred = np.array(blurred, dtype=np.uint8)
    return blurred

# 逆滤波函数
def inverse_filtering(input_img, PSF, eps=1e-5):
    # 获取输入图像的高度和宽度
    img_h, img_w = input_img.shape[:2]

    # 将模糊核填充到与输入图像相同的尺寸
    PSF_padded = np.zeros_like(input_img, dtype=np.float32)
    kh, kw = PSF.shape
    PSF_padded[:kh, :kw] = PSF  # 将PSF放置在左上角

    # 对输入图像和模糊核进行傅里叶变换
    input_fft = np.fft.fft2(input_img)
    PSF_fft = np.fft.fft2(PSF_padded) + eps  # 避免除零

    # 执行逆滤波：在频域上将输入图像除以模糊核的频谱
    result_fft = input_fft / PSF_fft

    # 通过逆傅里叶变换还原到空域
    result = np.abs(np.fft.ifft2(result_fft))

    # 将结果限制在 [0, 255] 范围内，并转换为 8 位整数
    result = np.clip(result, 0, 255)
    return result.astype(np.uint8)

# 加载图像
image_path = "img/kenny.jpg"  # 修改为你的实际路径
image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)  # 使用灰度A模式读取

# 生成模糊核并应用模糊
kernel_size = 15
angle = 60
PSE = make_PSE(kernel_size, angle)
blurred_image = motion_blur(image, PSE)

# 使用逆滤波复原
restored_image = inverse_filtering(blurred_image, PSE)

# 显示结果
plt.figure(figsize=(12, 4))
plt.subplot(1, 3, 1)
plt.title("Original Image")
plt.imshow(image, cmap='gray')
plt.axis('off')

plt.subplot(1, 3, 2)
plt.title("Blurred Image")
plt.imshow(blurred_image, cmap='gray')
plt.axis('off')

plt.subplot(1, 3, 3)
plt.title("Restored Image")
plt.imshow(restored_image, cmap='gray')
plt.axis('off')

plt.tight_layout()
plt.show()
